package com.heiye.blog.ai.model.openrouter;

import com.fasterxml.jackson.annotation.*;
import com.heiye.blog.framework.common.util.JsonUtils;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.model.ApiKey;
import org.springframework.ai.model.NoopApiKey;
import org.springframework.ai.model.SimpleApiKey;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.util.LinkedMultiValueMap;
import org.springframework.util.MultiValueMap;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.lang.reflect.Field;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.Consumer;
import java.util.function.Predicate;

/**
 * @author: heiye
 * @date: 2025/11/11 下午10:06
 * @version: v1.0.0
 * @description: openRouter api 设置
 */
@Slf4j
@Data
public class OpenRouterApi {
    /**
     * 地址前缀
     */
    private final String baseUrl;
    /**
     * api key
     */
    private final ApiKey apiKey;
    /**
     * webClient 结束
     */
    private static final Predicate<String> SSE_DONE_PREDICATE;
    /**
     * 通用地址
     */
    private final String completionsPath;
    /**
     * webClient 错误监听
     */
    private final ResponseErrorHandler responseErrorHandler;

    private final RestClient restClient;

    /**
     * 构建 webClient
     */
    private final WebClient webClient;

    /**
     * 构建 Builder
     *
     * @return
     */
    public static Builder builder() {
        return new Builder();
    }

    public OpenRouterApi(ApiKey apiKey, String baseUrl, MultiValueMap<String, String> headers, String completionsPath, ResponseErrorHandler responseErrorHandler, RestClient.Builder restClientBuilder, WebClient.Builder webClientBuilder) {
        this.apiKey = apiKey;
        this.baseUrl = baseUrl;
        this.completionsPath = completionsPath;
        this.responseErrorHandler = responseErrorHandler;

        // 创建请求头
        Consumer<HttpHeaders> finalHeaders = (h) -> {
            if (!(apiKey instanceof NoopApiKey)) {
                h.setBearerAuth(apiKey.getValue());
            }

            h.setContentType(MediaType.APPLICATION_JSON);
            h.addAll(headers);
        };

        this.restClient = restClientBuilder.baseUrl(baseUrl).defaultHeaders(finalHeaders).defaultStatusHandler(responseErrorHandler).build();
        this.webClient = webClientBuilder.baseUrl(baseUrl).defaultHeaders(finalHeaders).build();
    }

    public Flux<ChatCompletionChunk> chatCompletionStream(ChatCompletionRequest chatRequest) {
        return this.chatCompletionStream(chatRequest, new LinkedMultiValueMap());
    }

    /**
     * 调用 OpenRouter api
     *
     * @param chatRequest
     * @param additionalHttpHeader
     * @return
     */
    public Flux<ChatCompletionChunk> chatCompletionStream(ChatCompletionRequest chatRequest, MultiValueMap<String, String> additionalHttpHeader) {
        return this.webClient.post()
                // 请求地址
                .uri(this.getEndpoint(chatRequest))
                // 设置请求头
                .headers(headers -> headers.addAll(additionalHttpHeader))
                // 设置请求体
                .body(Mono.just(chatRequest), ChatCompletionRequest.class)
                // 获取响应
                .retrieve()
                // 获取响应体
                .bodyToFlux(String.class)
                // 到结束标记停止
                .takeUntil(SSE_DONE_PREDICATE)
                // 过滤掉结束标记
                .filter(SSE_DONE_PREDICATE.negate())
                // 转换为对象
                .map(content -> JsonUtils.parseObject(content, ChatCompletionChunk.class));
    }

    /**
     * 获取请求地址
     *
     * @param request
     * @return
     */
    private String getEndpoint(ChatCompletionRequest request) {
        // 这里可以根据调用的模型参数中是否带有前缀，重定向其他地址
        // 当前不需要删了
        return this.completionsPath;
    }

    static {
        SSE_DONE_PREDICATE = "[DONE]"::equals;
    }

    public static class Builder {
        private String baseUrl = "https://openrouter.ai/api/";
        private ApiKey apiKey;
        private MultiValueMap<String, String> headers = new LinkedMultiValueMap<>();
        private String completionsPath = "/v1/chat/completions";
        private RestClient.Builder restClientBuilder = RestClient.builder();
        private WebClient.Builder webClientBuilder = WebClient.builder();
        private ResponseErrorHandler responseErrorHandler;

        public Builder() {
            this.responseErrorHandler = RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER;
        }

        public Builder baseUrl(String baseUrl) {
            this.baseUrl = baseUrl;
            return this;
        }

        public Builder apiKey(ApiKey apiKey) {
            this.apiKey = apiKey;
            return this;
        }

        public Builder apiKey(String simpleApiKey) {
            this.apiKey = new SimpleApiKey(simpleApiKey);
            return this;
        }

        public Builder headers(MultiValueMap<String, String> headers) {
            this.headers = headers;
            return this;
        }

        public Builder completionsPath(String completionsPath) {
            this.completionsPath = completionsPath;
            return this;
        }

        public Builder restClientBuilder(RestClient.Builder restClientBuilder) {
            this.restClientBuilder = restClientBuilder;
            return this;
        }

        public Builder webClientBuilder(WebClient.Builder webClientBuilder) {
            this.webClientBuilder = webClientBuilder;
            return this;
        }

        public Builder responseErrorHandler(ResponseErrorHandler responseErrorHandler) {
            this.responseErrorHandler = responseErrorHandler;
            return this;
        }

        public OpenRouterApi build() {
            return new OpenRouterApi(this.apiKey, this.baseUrl, this.headers, this.completionsPath, this.responseErrorHandler, this.restClientBuilder, this.webClientBuilder);
        }
    }

    /**
     * OpenRouterRequest 请求记录类
     *
     * @param messages
     * @param prompt
     * @param model
     * @param responseFormat
     * @param stop
     * @param stream
     * @param maxTokens
     * @param temperature
     * @param seed
     * @param topP
     * @param topK
     * @param frequencyPenalty
     * @param presencePenalty
     * @param repetitionPenalty
     * @param logitBias
     * @param topLogprobs
     * @param minP
     * @param topA
     * @param transforms
     * @param models
     */
    @JsonInclude(JsonInclude.Include.NON_NULL)
    public record ChatCompletionRequest(
            // 聊天消息
            @JsonProperty("messages")
            List<ChatCompletionRequestMessage> messages,

            // 提示词
            @JsonProperty("prompt")
            String prompt,

            // 模型名称
            @JsonProperty("model")
            String model,

            // 结构化输出响应格式
            @JsonProperty("response_format")
            Object responseFormat,

            // 如果模型遇到停止数组中指定的任何标记，则立即停止生成。
            @JsonProperty("stop")
            List<String> stop,

            // 是否流式返回结果
            @JsonProperty("stream")
            Boolean stream,

            // 这设置了模型可以生成的标记数量的上限。
            // 它不会产生超过此限制的数量。
            // 最大值是上下文长度减去提示长度。
            // Range: [1, context_length)
            @JsonProperty("max_tokens")
            Integer maxTokens,

            // 此设置会影响模型响应的多样性。较低的值会导致更可预测和典型的响应，
            // 而较高的值则鼓励更多样化和不太常见的响应。
            // 在 0 时，模型始终对给定输入给出相同的响应。
            // Range: [0, 2]
            @JsonProperty("temperature")
            Double temperature,

            // 如果指定，推理将确定性地采样，
            // 以便具有相同种子和参数的重复请求应返回相同的结果。
            // 某些模型不能保证确定性。
            // Integer only
            @JsonProperty("seed")
            Integer seed,

            // 此设置将模型的选择限制为可能标记的百分比：只有概率加起来为 P 的顶级标记。
            // 较低的值使模型的响应更可预测，而默认设置允许全方位的标记选择。
            // 把它想象成一个动态的 Top-K。
            // Range: (0, 1]
            @JsonProperty("top_p")
            Double topP,

            // 这限制了模型在每个步骤中对标记的选择，使其从较小的集合中进行选择。
            // 值为 1 表示模型将始终选择最有可能的下一个标记，从而产生可预测的结果。
            // 默认情况下，此设置处于禁用状态，使模型考虑所有选择。
            // Range: [1, Infinity) Not available for OpenAI models
            @JsonProperty("top_k")
            Integer topK,

            // 此设置旨在根据标记在输入中出现的频率来控制标记的重复。
            // 它尝试降低使用那些在输入中出现较多的标记，与它们出现的频率成正比。
            // 令牌惩罚随出现次数而增加。负值将鼓励 token 重用。
            // Range: [-2, 2]
            @JsonProperty("frequency_penalty")
            Double frequencyPenalty,

            // 调整模型重复输入中已使用的特定标记的频率。
            // 较高的值使这种重复的可能性降低，而负值则相反。
            // 令牌惩罚不会随出现次数而变化。负值将鼓励代币重用。
            // Range: [-2, 2]
            @JsonProperty("presence_penalty")
            Double presencePenalty,

            // 有助于减少输入中标记的重复。较高的值使模型不太可能重复标记，
            // 但过高的值会使输出的连贯性降低（通常使用缺少小单词的连续句子）。
            // token 惩罚根据原始 token 的概率进行缩放。
            // Range: (0, 2]
            @JsonProperty("repetition_penalty")
            Double repetitionPenalty,

            // 接受一个 JSON 对象，该对象将令牌（由分词器中的令牌 ID 指定）映射到 -100 到 100 之间的关联偏差值。
            // 从数学上讲，偏差在采样之前被添加到模型生成的 logit 中。
            // 确切的效果因模型而异，但值介于 -1 和 1 之间应降低或增加选择的可能性;
            // 像 -100 或 100 这样的值应该会导致相关令牌被禁止或独占选择。
            // Range: [0, 1]
            @JsonProperty("logit_bias")
            Map<Integer, Integer> logitBias,

            // 一个介于 0 和 20 之间的整数，指定在每个标记位置最有可能返回的标记数，每个标记都有关联的对数概率。
            // 如果使用此参数，则必须将 logprobs 设置为 true。
            // Integer only
            @JsonProperty("top_logprobs")
            Integer topLogprobs,

            // 表示相对于最可能令牌的概率，考虑令牌的最小概率。
            // 该值根据最可能令牌的置信度而变化。如果您的 Min-P 设置为 0.1，
            // 则意味着它只允许至少是最佳选项的 1/10 的 token 。
            // Range: [0, 1]
            @JsonProperty("min_p")
            Double minP,

            // 根据最有可能的代币的概率，仅考虑具有“足够高”概率的顶级代币。可以把它想象成一个动态的 Top-P。
            // 较低的 Top-A 值会根据最高概率标记集中选择，但范围更窄。
            // 较高的 Top-A 值不一定会影响输出的创造力，而是根据最大概率细化过滤过程。
            // Range: [0, 1]
            @JsonProperty("top_a")
            Double topA,

            // 这对于不需要完美回忆的情况很有用。
            // 转换的工作原理是从提示中间删除或截断消息，直到提示适合模型的上下文窗口。
            @JsonProperty("transforms")
            List<String> transforms,

            // 允许您在主模型的提供程序关闭、速率受限或由于内容审核而拒绝回复时自动尝试其他模型。
            @JsonProperty("models")
            List<String> models,

            // 深度思考对象字段
            @JsonProperty("reasoning")
            Think reasoning) {
        public ChatCompletionRequest(List<ChatCompletionRequestMessage> messages, String model, Boolean stream, Double temperature) {
            this(messages, null, model, null, null, stream, null, temperature, null, null, null, null, null, null, null, null, null, null, null, null, null);
        }

        public ChatCompletionRequest(List<ChatCompletionRequestMessage> messages, Boolean stream) {
            this(messages, null, null, null, null, stream, null, null, null, null, null, null, null, null, null, null, null, null, null, null, null);
        }

        public ChatCompletionRequest(@JsonProperty("messages") List<ChatCompletionRequestMessage> messages, @JsonProperty("prompt") String prompt, @JsonProperty("model") String model, @JsonProperty("response_format") Object responseFormat, @JsonProperty("stop") List<String> stop, @JsonProperty("stream") Boolean stream, @JsonProperty("max_tokens") Integer maxTokens, @JsonProperty("temperature") Double temperature, @JsonProperty("seed") Integer seed, @JsonProperty("top_p") Double topP, @JsonProperty("top_k") Integer topK, @JsonProperty("frequency_penalty") Double frequencyPenalty, @JsonProperty("presence_penalty") Double presencePenalty, @JsonProperty("repetition_penalty") Double repetitionPenalty, @JsonProperty("logit_bias") Map<Integer, Integer> logitBias, @JsonProperty("top_logprobs") Integer topLogprobs, @JsonProperty("min_p") Double minP, @JsonProperty("top_a") Double topA, @JsonProperty("transforms") List<String> transforms, @JsonProperty("models") List<String> models, @JsonProperty("reasoning") Think reasoning) {
            this.messages = messages;
            this.prompt = prompt;
            this.model = model;
            this.responseFormat = responseFormat;
            this.stop = stop;
            this.stream = stream;
            this.maxTokens = maxTokens;
            this.temperature = temperature;
            this.seed = seed;
            this.topP = topP;
            this.topK = topK;
            this.frequencyPenalty = frequencyPenalty;
            this.presencePenalty = presencePenalty;
            this.repetitionPenalty = repetitionPenalty;
            this.logitBias = logitBias;
            this.topLogprobs = topLogprobs;
            this.minP = minP;
            this.topA = topA;
            this.transforms = transforms;
            this.models = models;
            this.reasoning = reasoning;
        }
    }

    @JsonInclude(JsonInclude.Include.NON_NULL)
    public record Think(
            // 分配多少 token 用于推理
            @JsonProperty("effort") Effort effort,
            // 推理的 token 数量
            @JsonProperty("max_tokens") Integer maxTokens,
            // 是否排除推理结果
            @JsonProperty("exclude") Boolean exclude,
            // 是否启用推理
            @JsonProperty("enabled") Boolean enabled) {

        public Think(@JsonProperty("effort") Effort effort, Integer maxTokens, Boolean exclude, Boolean enabled) {
            this.effort = effort;
            this.maxTokens = maxTokens;
            this.exclude = exclude;
            this.enabled = enabled;
        }
    }

    /**
     * 调用 openrount 返回响应数据
     */
    public record ChatCompletionChunk(
            @JsonProperty("id") String id,
            // AI 响应数据
            @JsonProperty("choices") List<Choice> choices,
            @JsonProperty("created") Long created,
            @JsonProperty("model") String model,

            @JsonProperty("object") String object,
            // 当提供商支持时才返回

            @JsonProperty("system_fingerprint") String systemFingerprint,
            // 对于非流式处理，始终返回使用情况数据。
            // 流式传输时，您将在末尾获得一个使用对象，
            // 并附有一个空的选择数组
            @JsonProperty("usage") Usage usage
    ) {
        /**
         * 调用 openrount 主要响应数据
         */
        public record Choice(

                @JsonProperty("index") Long index,

                // 模型生成的delta。存储 AI 的消息和推理过程
                @JsonProperty("delta") ChatCompletionResponseMessage delta,

                // 模型生成文本的完成原因。
                @JsonProperty("finish_reason") ChatCompletionFinishReason finishReason,

                // 模型生成文本的完成原因。
                // 有些模型不会提供 openrount 中 finishReason 里存在的字段，而是会用别的字段
                @JsonProperty("native_finish_reason") ChatCompletionFinishReason nativeFinishReason
        ) {
        }
    }

    /**
     * 记录 token 消耗
     */
    public record Usage(
            //  如果有图片和工具，则包含在内
            @JsonProperty("prompt_tokens") Integer promptTokens,
            // 生成的令牌
            @JsonProperty("completion_tokens") Integer completionTokens,
            // 上述两个字段的和
            @JsonProperty("total_tokens") Integer totalTokens
    ) {
    }

    /**
     * 创建请求文本消息
     *
     * @param role
     * @param content
     */
    @JsonInclude(JsonInclude.Include.NON_NULL)
    public record ChatCompletionRequestMessage(
            @JsonProperty("role") Role role,
            @JsonProperty("content") List<Content> content
    ) {
        /**
         * 创建文本消息
         *
         * @param role
         * @param text
         * @return
         */
        public static ChatCompletionRequestMessage OpenRouterChatCompletionTextMessage(Role role, String text) {
            return new ChatCompletionRequestMessage(role, Collections.singletonList(new TextContent(MessageType.TEXT, text)));
        }

        public ChatCompletionRequestMessage(@JsonProperty("role") Role role, @JsonProperty("content") List<Content> content) {
            this.role = role;
            this.content = content;
        }

        @JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type")
        @JsonSubTypes({
                @JsonSubTypes.Type(value = TextContent.class, name = "text"),
                @JsonSubTypes.Type(value = ImageContent.class, name = "image_url")
        })
        public interface Content {
        }

        // 文本内容
        @JsonInclude(JsonInclude.Include.NON_NULL)
        public record TextContent(
                @JsonProperty("type")
                MessageType type,
                @JsonProperty("text")
                String text
        ) implements Content {
        }

        // 图片内容
        @JsonInclude(JsonInclude.Include.NON_NULL)
        public record ImageContent(

                @JsonProperty("type")
                MessageType type,
                @JsonProperty("image_url")
                ImageUrl imageUrl
        ) implements Content {
            // 图片地址记录类
            public record ImageUrl(
                    // 图片地址 或 base64
                    @JsonProperty("url")
                    String url,
                    // 该参数默认 auto
                    @JsonProperty("detail")
                    String detail
            ) {
            }
        }
    }

    /**
     * 创建响应文本消息
     *
     * @param role
     * @param content
     * @param reasoning
     * @param reasoningDetails
     */
    @JsonInclude(JsonInclude.Include.NON_NULL)
    public record ChatCompletionResponseMessage(
            @JsonProperty("role") Role role,
            @JsonProperty("content") String content,
            // 深度思考推理过程
            @JsonProperty("reasoning") String reasoning,
            // 推理过程详情
            @JsonProperty("reasoning_details") List<OpenRouterReasoningDetail> reasoningDetails
    ) {


        /**
         * 推理过程详情
         */
        public record OpenRouterReasoningDetail(
                // 推理过程详情类型
                @JsonProperty("type") String type,
                // 推理过程详情内容
                @JsonProperty("text") String text,
                // 推理过程详情格式
                @JsonProperty("format") String format,
                // 推理过程详情索引
                @JsonProperty("index") Long index
        ) {
        }
    }

    // 角色枚举类
    public enum Role {
        @JsonProperty("system")
        SYSTEM,
        @JsonProperty("user")
        USER,
        @JsonProperty("assistant")
        ASSISTANT,
        @JsonProperty("tool")
        TOOL;

        private Role() {
        }
    }

    // 消息类型枚举类
    public enum MessageType {
        @JsonProperty("text")
        TEXT,
        @JsonProperty("image_url")
        IMAGE_URL
    }

    /**
     * 聊天完成原因枚举类
     */
    public enum ChatCompletionFinishReason {
        @JsonProperty("stop")
        STOP,
        @JsonProperty("length")
        LENGTH,
        @JsonProperty("content_filter")
        CONTENT_FILTER,
        @JsonProperty("tool_calls")
        TOOL_CALLS,
        @JsonProperty("tool_call")
        TOOL_CALL,
        @JsonProperty("error")
        ERROR;

        private ChatCompletionFinishReason() {
        }

        // 添加一个容错的方法
        @JsonCreator
        public static ChatCompletionFinishReason fromString(String value) {
            if (value == null) {
                return null;
            }

            // 统一转换为小写进行比较
            String normalized = value.toLowerCase();
            for (ChatCompletionFinishReason reason : values()) {
                if (reason.name().toLowerCase().equals(normalized)) {
                    return reason;
                }
            }
            // 如果都不匹配，返回 null 或默认值
            return null;
        }
    }

    /**
     * 创建深度思考参数类
     */
    public enum Effort {
        @JsonProperty("high")
        HIGH,
        @JsonProperty("medium")
        MEDIUM,
        @JsonProperty("low")
        LOW,
        ;

        private Effort() {
        }

        // 添加一个容错的方法
        @JsonCreator
        public static Effort fromString(String value) {
            if (value == null) {
                return null;
            }

            // 统一转换为小写进行比较
            String normalized = value.toLowerCase();
            for (Effort thinkEffort : values()) {
                if (thinkEffort.name().toLowerCase().equals(normalized)) {
                    return thinkEffort;
                }
            }
            // 如果都不匹配，返回 null 或默认值
            return null;
        }
    }
}
